import numpy as np
import matplotlib.pyplot as plt

from sklearn.metrics import classification_report
from sklearn.preprocessing import LabelBinarizer
# from tensorflow.keras.layers import Dense
# from tensorflow.keras.models import Sequential
from tensorflow.keras.optimizers import SGD
from tensorflow.keras.datasets import cifar10

from shallow_net import Shallow_Net


def main():
    print("[info]:开始加载CIFAR10数据集")
    (train_x, train_y), (test_x, test_y) = cifar10.load_data()
    # 归一化
    train_x = train_x.astype("float") / 255.0
    test_x = test_x.astype("float") / 255.0
    # 标签编码
    lb = LabelBinarizer()
    train_y = lb.fit_transform(train_y)
    test_y = lb.transform(test_y)
    label_names = ["airplane", "automobile", "bird", "cat", "deer",
                  "dog","frog","horse","ship","truck"]

    # 编译模型
    opt = SGD(0.01)
    model = Shallow_Net.build(width=32, height=32, depth=3, classes=10)
    model.compile(loss="categorical_crossentropy", optimizer=opt,
                  metrics=["acc"])
    # 训练
    record = model.fit(train_x, train_y, validation_data=(test_x, test_y),
              batch_size=32, epochs=40, verbose=1)

    # 评估模型
    predictions = model.predict(test_x, batch_size=32)
    print(classification_report(test_y.argmax(1),
                                predictions.argmax(1),
                                target_names=label_names))

    # 画图
    plt.style.use("ggplot")
    plt.figure()
    plt.plot(np.arange(0, 40), record.history["loss"], label="train_loss")
    plt.plot(np.arange(0, 40), record.history["val_loss"], label="val_loss")
    plt.plot(np.arange(0, 40), record.history["acc"], label="train_acc")
    plt.plot(np.arange(0, 40), record.history["val_acc"], label="val_acc")
    plt.title("Training Loss and Accuracy")
    plt.xlabel("Epoch #")
    plt.ylabel("Loss/Accuracy")
    plt.legend()
    plt.show()


if __name__ == '__main__':
    main()
